#!/bin/bash

# 定义变量

# train: /apdcephfs/share_976139/users/kaixunhuang/tools/env/megatron_omni_train
GPUS_PER_NODE=8  # 每个节点上的 GPU 数量
MASTER_PORT=12323  # 指定端口
DATA_PATH=/teaspeech_ceph/share_976139/users/kaixunhuang/data/chat/train_set/omni_train_2025_09_25_think_v6/train.jsonl
VALID_DATA_PATH=/teaspeech_ceph/share_976139/users/kaixunhuang/data/chat/train_set/omni_train_2025_09_25_think_v6/dev.jsonl
SAVE_PREFIX=qwen2_5omni_thinker_train_v0.31_xlgeng
SCRIPT="scripts/train_qwen2_5omni_thinker.sh"

. utils/parse_options.sh

echo $NODE_IP_LIST | awk -F'[,:]' '{for(i=1;i<=NF;i+=2) printf "%s slots=%s\n", $i, $(i+1)}' > ${SAVE_PREFIX}.hostfile
HOSTFILE=$SAVE_PREFIX.hostfile

# 解析hostfile，提取节点IP地址
NODES=($(awk '{print $1}' $HOSTFILE))
NNODES=${#NODES[@]}  # 计算总节点数
SLOTS=($(awk '{print $2}' $HOSTFILE | sed 's/slots=//g'))  # 提取slots信息并去掉' slots='前缀
# 获取主节点IP地址（用第一个节点作为master）
MASTER_ADDR=${NODES[0]}

trap "kill_slaves" SIGINT

kill_slaves() {
  echo "捕获到 Ctrl+C，终止从节点上的训练进程..."
  # 使用pdsh发送kill命令到所有从节点，终止训练进程
  for i in "${!NODES[@]}"; do
    NODE_IP=${NODES[$i]}
    echo "正在终止 ${NODE_IP} 上的进程..."
    pdsh -w $NODE_IP "ps -ef | grep finetune_audiollm| awk '{print $2}' | xargs kill -9 " &
  done
  wait
  echo "所有从节点进程已终止"
  exit 1
}

# 遍历所有节点并通过pdsh启动训练
for i in "${!NODES[@]}"; do
  NODE_RANK=$i  # 每个节点的rank值
  NODE_IP=${NODES[$i]}
  GPUS_PER_NODE=${SLOTS[$i]}
  # 计算世界规模
  WORLD_SIZE=$(($GPUS_PER_NODE * $NNODES))
  if [ $NODE_RANK -eq 0 ]; then

  pdsh -R ssh -w $NODE_IP "cd $PWD;bash $SCRIPT --MASTER_ADDR $MASTER_ADDR --MASTER_PORT $MASTER_PORT --NNODES $NNODES --NODE_RANK $NODE_RANK --WORLD_SIZE $WORLD_SIZE --DATA_PATH $DATA_PATH --VALID_DATA_PATH $VALID_DATA_PATH --SAVE_PREFIX $SAVE_PREFIX 2>&1 |tee logs/${SAVE_PREFIX}.$NODE_RANK.log" &
  #pdsh -R ssh -w $NODE_IP "cd $PWD;bash $SCRIPT --MASTER_ADDR $MASTER_ADDR --MASTER_PORT $MASTER_PORT --NNODES $NNODES --NODE_RANK $NODE_RANK --WORLD_SIZE $WORLD_SIZE --DATA_PATH $DATA_PATH --VALID_DATA_PATH $VALID_DATA_PATH --SAVE_PREFIX $SAVE_PREFIX |tee logs/${SAVE_PREFIX}.$NODE_RANK.log" &
  else
  # 使用pdsh在各个节点上运行pretrain.sh，并传递相关参数
  pdsh -R ssh -w $NODE_IP "cd $PWD;bash $SCRIPT --MASTER_ADDR $MASTER_ADDR --MASTER_PORT $MASTER_PORT --NNODES $NNODES --NODE_RANK $NODE_RANK --WORLD_SIZE $WORLD_SIZE --DATA_PATH $DATA_PATH --VALID_DATA_PATH $VALID_DATA_PATH --SAVE_PREFIX $SAVE_PREFIX |tee logs/${SAVE_PREFIX}.$NODE_RANK.log" &
  fi
done
wait
